#include <gtest/gtest.h>
#include <spdlog/spdlog.h>
#include <torch/torch.h>

#include <chrono>
#include <cmath>
#include <ctime>
#include <iomanip>
#include <limits>
#include <map>
#include <string>
#include <vector>

#include "hbm_abs_bw.h"
#include "wandbcpp.hpp"

/**
 * 测试hbm_abs_bw内核在不同数据大小下的性能
 * 使用wandb-cpp记录性能指标
 */
class HbmAbsBwPerformanceTest : public ::testing::Test {
 protected:
  void SetUp() override {
    // 初始化wandb - 使用更简单的形式
    wandbcpp::init({.project = "cuda_kernels_performance",
                    .tags = {"performance", "hbm_abs_bw"}});

    // 获取GPU信息并记录到日志
    int device;
    cudaGetDevice(&device);
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, device);

    spdlog::info("GPU: {}", prop.name);
    spdlog::info("Compute Capability: {}.{}", prop.major, prop.minor);

    // 获取可用GPU内存
    size_t free_mem, total_mem;
    cudaMemGetInfo(&free_mem, &total_mem);
    spdlog::info("GPU Memory: Free: {:.2f} GB, Total: {:.2f} GB",
                 free_mem / 1e9, total_mem / 1e9);
    max_available_memory_ = free_mem * 0.8;  // 使用80%的可用内存作为上限
  }

  void TearDown() override {
    // 结束wandb会话
    wandbcpp::finish();
  }

  /**
   * 测试特定配置下的性能
   *
   * @tparam VectorSize - 向量大小 (1, 2, 或 4)
   * @tparam Unroll - 展开因子 (1, 2, 或 4)
   * @param n - 元素数量
   * @return 带宽 (GB/s)
   */
  template <int VectorSize, int Unroll>
  float TestPerformance(size_t n) {
    // 检查内存需求是否超过可用内存
    size_t memory_required = n * sizeof(float) * 2;  // 输入和输出
    if (memory_required > max_available_memory_) {
      spdlog::warn(
          "Skipping test for n={} as it requires {:.2f} GB which exceeds "
          "available memory",
          n, memory_required / 1e9);
      return 0.0f;
    }

    // 确保n是VectorSize的倍数
    size_t aligned_n = (n / VectorSize) * VectorSize;
    if (aligned_n != n) {
      spdlog::warn(
          "Adjusting n from {} to {} to ensure it's a multiple of "
          "VectorSize={}",
          n, aligned_n, VectorSize);
      n = aligned_n;
    }

    // 创建输入张量
    torch::Tensor input_tensor;
    torch::Tensor output_tensor;

    try {
      input_tensor =
          torch::randn({static_cast<int64_t>(n)},
                       torch::TensorOptions().dtype(torch::kFloat32));

      // 创建输出张量
      output_tensor = torch::zeros_like(input_tensor);
    } catch (const std::exception& e) {
      spdlog::error("Failed to allocate CPU tensors: {}", e.what());
      return 0.0f;
    }

    // 将数据移动到GPU
    try {
      input_cuda = input_tensor.cuda();
      output_cuda = output_tensor.cuda();
    } catch (const std::exception& e) {
      spdlog::error("Failed to move tensors to GPU: {}", e.what());
      return 0.0f;
    }

    // 创建CUDA流
    cudaStream_t stream;
    cudaError_t err = cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
    if (err != cudaSuccess) {
      spdlog::error("Failed to create CUDA stream: {}",
                    cudaGetErrorString(err));
      return 0.0f;
    }

    // 创建CUDA事件来测量时间
    cudaEvent_t start, stop;
    err = cudaEventCreate(&start);
    if (err != cudaSuccess) {
      spdlog::error("Failed to create CUDA event: {}", cudaGetErrorString(err));
      cudaStreamDestroy(stream);
      return 0.0f;
    }

    err = cudaEventCreate(&stop);
    if (err != cudaSuccess) {
      spdlog::error("Failed to create CUDA event: {}", cudaGetErrorString(err));
      cudaEventDestroy(start);
      cudaStreamDestroy(stream);
      return 0.0f;
    }

    // 预热
    for (int i = 0; i < 5; i++) {
      hbm::pointwise::launch_abs_kernel<VectorSize, Unroll>(
          static_cast<float*>(output_cuda.data_ptr()),
          static_cast<const float*>(input_cuda.data_ptr()), n, stream);

      // 检查是否有错误
      err = cudaGetLastError();
      if (err != cudaSuccess) {
        spdlog::error("CUDA error during warmup: {}", cudaGetErrorString(err));
        cudaEventDestroy(start);
        cudaEventDestroy(stop);
        cudaStreamDestroy(stream);
        return 0.0f;
      }
    }

    err = cudaStreamSynchronize(stream);
    if (err != cudaSuccess) {
      spdlog::error("CUDA error synchronizing stream after warmup: {}",
                    cudaGetErrorString(err));
      cudaEventDestroy(start);
      cudaEventDestroy(stop);
      cudaStreamDestroy(stream);
      return 0.0f;
    }

    // 开始计时
    cudaEventRecord(start, stream);

    // 运行内核10次，记录总时间
    const int repeat_count = 10;
    for (int i = 0; i < repeat_count; i++) {
      hbm::pointwise::launch_abs_kernel<VectorSize, Unroll>(
          static_cast<float*>(output_cuda.data_ptr()),
          static_cast<const float*>(input_cuda.data_ptr()), n, stream);

      // 检查是否有错误
      err = cudaGetLastError();
      if (err != cudaSuccess) {
        spdlog::error("CUDA error during benchmark: {}",
                      cudaGetErrorString(err));
        cudaEventDestroy(start);
        cudaEventDestroy(stop);
        cudaStreamDestroy(stream);
        return 0.0f;
      }
    }

    // 结束计时
    cudaEventRecord(stop, stream);
    err = cudaStreamSynchronize(stream);
    if (err != cudaSuccess) {
      spdlog::error("CUDA error synchronizing stream after benchmark: {}",
                    cudaGetErrorString(err));
      cudaEventDestroy(start);
      cudaEventDestroy(stop);
      cudaStreamDestroy(stream);
      return 0.0f;
    }

    // 计算耗时
    float total_milliseconds = 0;
    cudaEventElapsedTime(&total_milliseconds, start, stop);

    // 计算平均每次运行的时间
    float avg_milliseconds = total_milliseconds / repeat_count;

    // 计算带宽 (GB/s)
    float bytes_processed = n * sizeof(float) * 2;  // 读和写
    float bandwidth = bytes_processed / (avg_milliseconds * 1e-3) / 1e9;

    // 清理
    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    cudaStreamDestroy(stream);

    return bandwidth;
  }

  /**
   * 运行所有配置的性能测试
   *
   * @param n - 元素数量
   */
  void RunAllConfigurations(size_t n) {
    spdlog::info("Testing performance with n = {}", n);

    // 设置重复次数常量
    const int repeat_count = 10;
    spdlog::info("Running each kernel {} times and taking average",
                 repeat_count);

    // 测试所有配置
    float bw_1_1 = TestPerformance<1, 1>(n);
    if (bw_1_1 == 0.0f) {
      spdlog::warn("Skipping remaining tests for n={} due to previous errors",
                   n);
      return;
    }

    float bw_1_2 = TestPerformance<1, 2>(n);
    float bw_1_4 = TestPerformance<1, 4>(n);
    float bw_2_1 = TestPerformance<2, 1>(n);
    float bw_2_2 = TestPerformance<2, 2>(n);
    float bw_2_4 = TestPerformance<2, 4>(n);
    float bw_4_1 = TestPerformance<4, 1>(n);
    float bw_4_2 = TestPerformance<4, 2>(n);
    float bw_4_4 = TestPerformance<4, 4>(n);

    // 计算数据大小
    double data_size_mb =
        static_cast<double>(n * sizeof(float)) / (1024 * 1024);

    // 记录到wandb - 简化版本，只记录带宽
    wandbcpp::log({
        {"data_size", static_cast<double>(n)},
        {"data_size_mb", data_size_mb},
        {"vs1_u1", bw_1_1},
        {"vs1_u2", bw_1_2},
        {"vs1_u4", bw_1_4},
        {"vs2_u1", bw_2_1},
        {"vs2_u2", bw_2_2},
        {"vs2_u4", bw_2_4},
        {"vs4_u1", bw_4_1},
        {"vs4_u2", bw_4_2},
        {"vs4_u4", bw_4_4},
    });

    // 输出结果
    spdlog::info("Data Size: {} elements ({:.2f} MB)", n, data_size_mb);
    spdlog::info("VectorSize=1, Unroll=1: {:.2f} GB/s", bw_1_1);
    spdlog::info("VectorSize=1, Unroll=2: {:.2f} GB/s", bw_1_2);
    spdlog::info("VectorSize=1, Unroll=4: {:.2f} GB/s", bw_1_4);
    spdlog::info("VectorSize=2, Unroll=1: {:.2f} GB/s", bw_2_1);
    spdlog::info("VectorSize=2, Unroll=2: {:.2f} GB/s", bw_2_2);
    spdlog::info("VectorSize=2, Unroll=4: {:.2f} GB/s", bw_2_4);
    spdlog::info("VectorSize=4, Unroll=1: {:.2f} GB/s", bw_4_1);
    spdlog::info("VectorSize=4, Unroll=2: {:.2f} GB/s", bw_4_2);
    spdlog::info("VectorSize=4, Unroll=4: {:.2f} GB/s", bw_4_4);

    // 清理CUDA缓存，防止内存泄漏
    cudaDeviceSynchronize();
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
      spdlog::error("CUDA error after tests: {}", cudaGetErrorString(err));
    }

    // 使用PyTorch的内存释放机制
    // 让输入和输出张量超出作用域自动释放
    input_cuda = torch::Tensor();
    output_cuda = torch::Tensor();

    // 强制进行垃圾回收
    cudaDeviceSynchronize();
  }

 private:
  size_t max_available_memory_;
  torch::Tensor input_cuda;   // 添加成员变量以便清理
  torch::Tensor output_cuda;  // 添加成员变量以便清理
};

// 测试不同数据大小下的性能
TEST_F(HbmAbsBwPerformanceTest, DataSizeScaling) {
  // 从1K到1G，按2的倍数递增
  for (size_t n = 1024; n <= (1ULL << 30); n *= 2) {
    RunAllConfigurations(n);

    // 对于非常大的数据集，可能需要清理GPU内存
    if (n >= (1ULL << 28)) {  // 256M及以上
      cudaDeviceSynchronize();
      // 使用标准CUDA API清理内存
      cudaFree(0);  // 触发CUDA运行时清理缓存
    }
  }
}

// 测试少量数据大小（用于快速验证功能）
TEST_F(HbmAbsBwPerformanceTest, SmallTest) {
  // 只测试几个数据大小
  std::vector<size_t> sizes = {
      1024,               // 1K
      1024 * 1024,        // 1M
      10 * 1024 * 1024,   // 10M
      100 * 1024 * 1024,  // 100M
  };

  for (size_t n : sizes) {
    RunAllConfigurations(n);

    // 每次测试后清理缓存
    cudaDeviceSynchronize();
    // 使用标准CUDA API清理内存
    cudaFree(0);  // 触发CUDA运行时清理缓存
  }
}
